import os
import numpy as np
import math
import ast
import json
import pygame
import imageio
from datetime import datetime
from openai import OpenAI
from apikey import api_key, base_url, model_name
import tempfile
import shutil
import sys
from PIL import Image, ImageDraw, ImageFont


class LogicAgent:
    def __init__(self, env):
        self.env = env
        self.api_key = api_key
        self.base_url = base_url
        self.model_name = model_name
        self.current_plan = []
        self.current_task_index = 0  # 当前执行的任务索引
        self.current_target = None  # 当前任务的目标信息
        self.task_completed = False  # 当前任务是否完成
        self.moving_to_target = False  # 是否正在移动中
        self.target_joint_angles = None  # 存储目标关节角度

        # 创建OpenAI客户端
        self.client = OpenAI(
            api_key=self.api_key,
            base_url=self.base_url,
        )

    def generate_plan(self, task_description):
        """调用Qwen3模型生成任务规划"""
        # 构建场景描述
        scene_description = self._build_scene_description()

        # 构建LLM提示
        prompt = self._build_llm_prompt(scene_description, task_description)

        print("-" * 20)
        print(prompt)
        print("-" * 20)

        # 调用Qwen3 API
        plan = self._call_qwen3_api(prompt)

        print("-" * 20)
        print(plan)
        print("-" * 20)

        # 解析规划
        self.current_plan = self._parse_plan(plan)
        self.current_task_index = 0
        self.task_completed = False
        self.moving_to_target = False
        return self.current_plan

    def get_next_action(self):
        """获取下一个动作（关节角变化量）"""
        # 所有任务已完成
        if self.current_task_index >= len(self.current_plan):
            return None

        # 如果当前任务已完成，移动到下一个任务
        if self.task_completed:
            print("任务已完成，移动到下一个任务")
            self.current_task_index += 1
            self.task_completed = False
            self.moving_to_target = False
            self.target_joint_angles = None  # 重置目标关节角度

            # 检查是否还有任务
            if self.current_task_index >= len(self.current_plan):
                return None

        # 获取当前任务
        current_task = self.current_plan[self.current_task_index]
        action_type = current_task[0]

        # 处理移动任务
        if action_type == "move":
            # 如果是第一次进入移动任务，设置目标
            if not self.moving_to_target:
                # 设置目标位置和旋转
                self.current_target = {
                    "x": current_task[1],
                    "y": current_task[2],
                    "rotation": current_task[3],
                }
                self.moving_to_target = True

                # 计算逆向运动学
                self.target_joint_angles = self._inverse_kinematics(
                    self.current_target["x"],
                    self.current_target["y"],
                    self.current_target["rotation"],
                )
                print("逆向运动学结果：", self.target_joint_angles)

            # 检查是否到达目标位置
            if self._check_reached_target():
                self.task_completed = True
                return np.array([0, 0, 0, 0])  # 无动作
            else:
                # 计算关节角变化量
                delta_angles = self._calculate_delta_angles(self.target_joint_angles)
                # 返回关节角变化量
                return np.append(delta_angles, 0)  # 无夹爪动作

        # 处理抓取任务
        elif action_type == "catch":
            self.task_completed = True
            return np.array([0, 0, 0, 1])  # 抓取动作

        # 处理释放任务
        elif action_type == "release":
            self.task_completed = True
            return np.array([0, 0, 0, 2])  # 释放动作

        return None

    def _check_reached_target(self):
        """检查机械臂是否到达目标位置"""
        # 获取当前末端位姿
        current_pos = self.env._forward_kinematics()
        current_rotation = sum(self.env.joint_angles)

        # 计算位置误差
        pos_error = np.linalg.norm(
            np.array(
                [
                    self.current_target["x"] - current_pos[0],
                    self.current_target["y"] - current_pos[1],
                ]
            )
        )

        # 计算旋转误差
        rot_error = abs(
            self._wrap_angle(self.current_target["rotation"] - current_rotation)
        )

        # 检查是否在允许的误差范围内
        position_threshold = 0.05  # 5厘米
        rotation_threshold = 0.087  # 约5度

        return pos_error < position_threshold and rot_error < rotation_threshold

    def _build_scene_description(self):
        """构建场景描述文本"""
        desc = f"现在的场景是：\n"
        desc += f"环境尺寸为 ({self.env.WORLD_WIDTH}, {self.env.WORLD_HEIGHT})\n"

        # 获取机械臂末端位姿
        end_pos = self.env._forward_kinematics()
        rotation = sum(self.env.joint_angles)  # 总旋转角度
        desc += f"机械臂当前的末端位姿为 ({end_pos[0]:.2f}, {end_pos[1]:.2f}, {rotation:.2f})\n"

        # 添加物体描述
        objects = []
        for table in self.env.tables:
            objects.append(
                {
                    "type": "table",
                    "x": table["pos"][0],
                    "y": table["pos"][1],
                    "rotation": 0,
                    "color": self._get_color_name(table["color"]),
                }
            )

        for box in self.env.boxes:
            objects.append(
                {
                    "type": "box",
                    "x": box["pos"][0],
                    "y": box["pos"][1],
                    "rotation": 0,
                    "color": self._get_color_name(box["color"]),
                }
            )

        for obj in objects:
            desc += (
                f"在位置 ({obj['x']:.2f}, {obj['y']:.2f}) "
                f"有一个姿态为 {obj['rotation']} "
                f"类型为 {obj['type']} "
                f"的对象，颜色为 {obj['color']}。\n"
            )

        return desc

    def _build_llm_prompt(self, scene_desc, task_desc):
        """构建LLM提示，添加墙壁、天花板和地面的坐标描述"""
        # 添加环境边界描述
        boundaries_desc = (
            f"环境边界定义如下：\n"
            f"- 地面：y = {self.env.WORLD_HEIGHT}（环境底部）\n"
            f"- 天花板：y = 0（环境顶部）\n"
            f"- 左侧墙壁：x = 0\n"
            f"- 右侧墙壁：x = {self.env.WORLD_WIDTH}\n"
            f"- 后侧墙壁：z = 0（在2D环境中忽略）\n"
            f"- 前侧墙壁：z = {self.env.WORLD_WIDTH}（在2D环境中忽略）\n"
        )

        prompt = (
            "请将任务分解为机械臂的动作序列，使用以下格式：\n"
            "[[动作类型, 目标x, 目标y, 目标旋转]]\n"
            "动作类型只能是 'move', 'catch', 'release' 三种：\n"
            "- 'move': 机械臂移动到指定位置 (目标x, 目标y, 目标旋转)\n"
            "- 'catch': 抓取对象，位姿参数无效\n"
            "- 'release': 释放对象，位姿参数无效\n"
            "输出示例：\n"
            '[["move", 1.5, 2.3, 0.78],\n'
            ' ["catch", 0, 0, 0],\n'
            ' ["move", 3.0, 1.0, 0.0],\n'
            ' ["release", 0, 0, 0]]\n'
            "坐标系定义为y越小越靠上；"
            "规划结束时机械臂末端必须处于释放状态;"
            "规划时注意，箱子之间不会发生碰撞，但是箱子和桌子之间会放生碰撞，所以注意处理中间过程;"
            "另外尽量保证使用少的操作，比如把一个箱子移动到它本来就来的位置这种操作不要进行输出;"
            "请直接输出动作序列，不要包含任何额外文本。\n"
            "下面是正式任务"
            f"{boundaries_desc}\n"  # 添加边界描述
            f"{scene_desc}\n"
            f"任务是：\n{task_desc}\n\n"
        )
        return prompt

    def _call_qwen3_api(self, prompt):
        """使用OpenAI风格的API调用Qwen3模型"""
        try:
            # 创建聊天完成请求
            completion = self.client.chat.completions.create(
                model=self.model_name,
                messages=[
                    {
                        "role": "system",
                        "content": "你是一个机械臂控制专家，负责将自然语言任务分解为机械臂动作序列。",
                    },
                    {"role": "user", "content": prompt},
                ],
                temperature=0.1,
                max_tokens=1000,
                # 对于Qwen3开源版模型，需要设置enable_thinking=False
                extra_body={"enable_thinking": False},
            )

            # 提取模型响应
            response = completion.choices[0].message.content
            return response

        except Exception as e:
            print(f"API调用异常: {e}")
            return ""

    def _parse_plan(self, plan_str):
        """解析规划字符串"""
        try:
            # 尝试直接解析JSON
            return json.loads(plan_str)
        except:
            try:
                # 尝试解析类似Python列表的字符串
                return ast.literal_eval(plan_str)
            except:
                # 尝试提取有效部分
                start = plan_str.find("[")
                end = plan_str.rfind("]") + 1
                if start != -1 and end != -1 and end > start:
                    return ast.literal_eval(plan_str[start:end])
                else:
                    print("无法解析规划:", plan_str)
                    return []

    def _inverse_kinematics(self, x, y, rotation):
        """
        逆向运动学：使用几何方法计算目标位姿对应的关节角度
        """
        print(f"正在进行移动到{(x,y,rotation)}的逆向运动学计算")
        # 基座位置
        base_x, base_y = self.env.base_pos

        # 连杆长度
        l1, l2, l3 = self.env.LINK_LENGTHS

        # 计算末端执行器的实际位置（考虑基座偏移）
        target_x = x - base_x
        target_y = y - base_y

        # 计算第三个关节的位置（末端位置减去第三个连杆在目标方向上的分量）
        # 注意：第三个连杆的方向由目标旋转决定
        joint2_x = target_x - l3 * math.cos(rotation)
        joint2_y = target_y - l3 * math.sin(rotation)

        # 计算第二个关节到基座的距离
        distance_to_joint2 = math.sqrt(joint2_x**2 + joint2_y**2)

        # 检查是否可达
        if distance_to_joint2 > l1 + l2 or distance_to_joint2 < abs(l1 - l2):
            # 如果不可达，调整到最近可达点
            if distance_to_joint2 > l1 + l2:
                scale = (l1 + l2) / distance_to_joint2
            else:
                scale = abs(l1 - l2) / distance_to_joint2
            joint2_x *= scale
            joint2_y *= scale
            print(f"目标位置不可达，已调整到最近点: ({joint2_x}, {joint2_y})")

        # 计算第一个关节角度 q1
        # 从基座到第二个关节的向量角度
        alpha = math.atan2(joint2_y, joint2_x)

        # 使用余弦定理计算第二个关节角度 q2
        cos_q2 = (joint2_x**2 + joint2_y**2 - l1**2 - l2**2) / (2 * l1 * l2)
        cos_q2 = np.clip(cos_q2, -1.0, 1.0)  # 确保在有效范围内
        q2 = math.acos(cos_q2)

        # 计算第一个关节角度 q1
        # 使用两个可能的解（肘部向上和肘部向下）
        # 这里我们选择肘部向下的解（q2为负）
        beta = math.atan2(l2 * math.sin(q2), l1 + l2 * math.cos(q2))
        q1 = alpha - beta

        # 计算第三个关节角度 q3
        q3 = rotation - q1 - q2

        # 将角度限制在[-π, π]范围内
        q1 = self._wrap_angle(q1)
        q2 = self._wrap_angle(q2)
        q3 = self._wrap_angle(q3)

        print(f"逆向运动学计算完成: q1={q1:.4f}, q2={q2:.4f}, q3={q3:.4f}")
        return np.array([q1, q2, q3])

    def _calculate_delta_angles(self, target_angles):
        """
        计算关节角度变化量
        """
        delta = []
        for i in range(3):
            current = self.env.joint_angles[i]
            target = target_angles[i]

            # 计算最小角度差
            diff = target - current
            diff = self._wrap_angle(diff)

            # 限制最大变化量
            max_step = 0.1  # 弧度
            if abs(diff) > max_step:
                diff = max_step if diff > 0 else -max_step

            delta.append(diff)

        return np.array(delta)

    def _wrap_angle(self, angle):
        """
        将角度限制在[-π, π]范围内
        """
        while angle > np.pi:
            angle -= 2 * np.pi
        while angle < -np.pi:
            angle += 2 * np.pi
        return angle

    def _get_color_name(self, rgb):
        """
        将RGB颜色转换为名称
        """
        colors = {
            tuple(self.env.BLACK): "黑色",
            tuple(self.env.WHITE): "白色",
            tuple(self.env.GRAY): "灰色",
            tuple(self.env.RED): "红色",
            tuple(self.env.GREEN): "绿色",
            tuple(self.env.BLUE): "蓝色",
            tuple(self.env.YELLOW): "黄色",
            tuple(self.env.PURPLE): "紫色",
            tuple(self.env.PINK): "粉色",
            tuple(self.env.ORANGE): "橙色",
        }
        return colors.get(tuple(rgb), "未知颜色")


def add_text_to_image(
    image_array, text, position=(10, 10), font_size=20, color=(0, 0, 0)
):
    """
    在图像上添加文字
    """
    # 将numpy数组转换为PIL图像
    image = Image.fromarray(image_array)

    # 创建一个可绘制的对象
    draw = ImageDraw.Draw(image)

    try:
        # 尝试加载支持中文的字体
        # Windows 系统常用字体
        font_paths = [
            "C:/Windows/Fonts/simhei.ttf",  # 黑体
            "C:/Windows/Fonts/simsun.ttc",  # 宋体
            "C:/Windows/Fonts/msyh.ttc",  # 微软雅黑
            "C:/Windows/Fonts/msyhbd.ttc",  # 微软雅黑粗体
        ]

        # macOS 系统常用字体
        mac_font_paths = [
            "/System/Library/Fonts/PingFang.ttc",  # 苹方
            "/System/Library/Fonts/STHeiti Light.ttc",  # 华文黑体
            "/System/Library/Fonts/STHeiti Medium.ttc",  # 华文黑体
        ]

        # Linux 系统常用字体
        linux_font_paths = [
            "/usr/share/fonts/truetype/droid/DroidSansFallbackFull.ttf",  # Droid Sans
            "/usr/share/fonts/truetype/wqy/wqy-microhei.ttc",  # 文泉驿微米黑
        ]

        # 尝试加载所有可能的字体路径
        font = None
        for path in font_paths + mac_font_paths + linux_font_paths:
            try:
                font = ImageFont.truetype(path, font_size)
                break  # 找到可用字体就停止
            except IOError:
                continue

        # 如果没有找到任何字体，使用默认字体
        if font is None:
            font = ImageFont.load_default()
            print("警告：未找到中文字体，使用默认字体（可能不支持中文）")

    except Exception as e:
        print(f"字体加载错误: {e}")
        font = ImageFont.load_default()

    # 添加文字
    draw.text(position, text, fill=color, font=font)

    # 将PIL图像转换回numpy数组
    return np.array(image)


# 测试代码
if __name__ == "__main__":
    from arm_env import ArmEnv
    import tempfile
    import shutil

    # 创建环境
    env = ArmEnv()
    env.reset()

    # 创建智能体
    agent = LogicAgent(env)

    # 创建临时文件夹保存帧
    temp_dir = tempfile.mkdtemp()
    frames = []
    frame_count = 0

    # 定义任务
    env.render()
    task_description = input("请输入任务描述: ")
    user_task = task_description  # 保存用户输入的指令

    # 生成规划
    plan = agent.generate_plan(task_description)
    print("生成的规划:")
    for i, step in enumerate(plan):
        print(f"步骤 {i+1}: {step}")

    # 保存初始状态
    pygame.image.save(
        env.screen, os.path.join(temp_dir, f"frame_{frame_count:04d}.png")
    )
    frames.append(
        imageio.imread(os.path.join(temp_dir, f"frame_{frame_count:04d}.png"))
    )
    frame_count += 1

    # 执行规划
    running = True
    step_count = 0  # 添加步骤计数器
    max_steps = 10000  # 最大步骤限制

    while running:
        step_count += 1  # 增加步骤计数

        # 检查是否超过最大步骤
        if step_count > max_steps:
            print(f"\n解算失败：超过最大步骤限制({max_steps})")
            print("可能原因：")
            print("1. 规划任务过于复杂")
            print("2. 机械臂无法到达目标位置")
            print("3. 环境中有障碍物阻挡")
            print("请尝试简化任务或调整环境")
            break  # 退出循环

        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False

        action = agent.get_next_action()
        print(f"步骤 {step_count} - 获得动作：", action)
        if action is None:
            print("规划执行完成")
            running = False
            continue

        # 如果动作不为空，执行动作
        if np.any(action[:3] != 0) or action[3] != 0:
            env.step(action)
            env.render()

            # 保存当前帧并立即读取到内存
            pygame.image.save(
                env.screen, os.path.join(temp_dir, f"frame_{frame_count:04d}.png")
            )
            frames.append(
                imageio.imread(os.path.join(temp_dir, f"frame_{frame_count:04d}.png"))
            )
            frame_count += 1
        else:
            # 空动作，复制前一帧
            frames.append(frames[-1])
            frame_count += 1

    # 创建GIF
    gif_filename = f"simulation_{datetime.now().strftime('%Y%m%d_%H%M%S')}.gif"

    # 在所有帧上添加用户指令
    frames_with_text = []
    for frame in frames:
        # 在左上角添加黑色文字
        frame_with_text = add_text_to_image(
            frame,
            f"任务: {user_task}",
            position=(10, 10),
            font_size=20,
            color=(0, 0, 0),  # 黑色文字
        )
        frames_with_text.append(frame_with_text)

    # 保存为GIF
    if frames_with_text:  # 确保有帧可保存
        imageio.mimsave(gif_filename, frames_with_text, fps=10, loop=0)
        print(f"仿真结果已保存为: {gif_filename}")
    else:
        print("没有生成仿真结果")

    # 清理临时文件
    shutil.rmtree(temp_dir)

    pygame.quit()

    # 如果超过最大步骤，退出程序
    if step_count > max_steps:
        sys.exit(1)  # 非正常退出
